#include "ECS/System/PUnlitMaterialSystem.h"
#include "ECS/Component/PTransformComponent.h"

#include "PiccoloFileUtil.h"
#include "PiccoloApplication.h"

#include "Graphics/PiccoloVKPipeline.h"
#include "Graphics/PiccoloVKDescriptorSet.h"
#include "Graphics/PiccoloVKImageView.h"
#include "Graphics/PiccoloVKFrameBuffer.h"

#include "Render/PiccoloRenderTarget.h"


namespace Piccolo
{
    void PUnlitMaterialSystem::OnInit(PiccoloVKRenderPass *renderPass)
    {
        PiccoloVKDevice* device = GetDevice();
        // Frame UBO
        {
            const std::vector<VkDescriptorSetLayoutBinding> bindings =
            {
                {
                    .binding = 0,
                    .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                    .descriptorCount = 1,
                    .stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT
                }
            };
            mFrameUBODescriptorSetLayout = std::make_shared<PiccoloVKDescriptorSetLayout>(device, bindings);
        }

        // Material Params
        {
            const std::vector<VkDescriptorSetLayoutBinding> bindings =
            {
                {
                    .binding = 0,
                    .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                    .descriptorCount = 1,
                    .stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT
                }
            };
            mMaterialParamDescriptorSetLayout = std::make_shared<PiccoloVKDescriptorSetLayout>(device, bindings);
        }

        // Material Resource
        {
            const std::vector<VkDescriptorSetLayoutBinding> bindings =
            {
                {
                    .binding = 0,
                    .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                    .descriptorCount = 1,
                    .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT
                },
                {
                    .binding = 1,
                    .descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                    .descriptorCount = 1,
                    .stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT
                }
            };
            mMaterialResourceDescriptorSetLayout = std::make_shared<PiccoloVKDescriptorSetLayout>(device, bindings);
        }

        // 准备常量和UBO参数给Pipeline Layout
        VkPushConstantRange modelPushConstant =
        {
            .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
            .offset = 0,
            .size = sizeof(ModelPushConstant)
        };
        ShaderLayout shaderLayout =
        {
            .descriptorSetLayouts = { mFrameUBODescriptorSetLayout->GetDescriptorSetLayout(), mMaterialParamDescriptorSetLayout->GetDescriptorSetLayout(), mMaterialResourceDescriptorSetLayout->GetDescriptorSetLayout() },
            .pushConstantRanges = { modelPushConstant }
        };
        std::vector<VkVertexInputBindingDescription> vertexBindings =
        {
            {
                .binding = 0,
                .stride = sizeof(PiccoloVertex),
                .inputRate = VK_VERTEX_INPUT_RATE_VERTEX
            }
        };
        std::vector<VkVertexInputAttributeDescription> vertexAttributes =
        {
            {
                .location = 0,
                .binding = 0,
                .format = VK_FORMAT_R32G32B32_SFLOAT,
                .offset = offsetof(PiccoloVertex, position)
            },
            {
                .location = 1,
                .binding = 0,
                .format = VK_FORMAT_R32G32_SFLOAT,
                .offset = offsetof(PiccoloVertex, texcoord0)
            },
            {
                .location = 2,
                .binding = 0,
                .format = VK_FORMAT_R32G32B32_SFLOAT,
                .offset = offsetof(PiccoloVertex, normal)
            }
        };

        mPipelineLayout = std::make_shared<PiccoloVKPipelineLayout>(device, PICCOLO_RES_SHADER_DIR"unlit_material.vert", PICCOLO_RES_SHADER_DIR"unlit_material.frag", shaderLayout);
        // 创建Pipeline
        mPipeline = std::make_shared<PiccoloVKPipeline>(device, renderPass, mPipelineLayout.get());
        mPipeline->SetVertexInputState(vertexBindings, vertexAttributes);
        mPipeline->EnableDepthTest();
        mPipeline->SetDynamicState({ VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR });
        mPipeline->SetMultisampleState(VK_SAMPLE_COUNT_4_BIT, VK_FALSE);
        mPipeline->Create();

        // 根据UBO的数量来创建指定数量的描述符池
        std::vector<VkDescriptorPoolSize> poolSizes =
        {
            {
                .type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                .descriptorCount = 1
            }
        };
        mDescriptorPool = std::make_shared<PiccoloVKDescriptorPool>(device, 1, poolSizes);
        mFrameUBODescriptorSet = mDescriptorPool->AllocateDescriptorSet(mFrameUBODescriptorSetLayout.get(), 1)[0];
        mFrameUBOBuffer = std::make_shared<PiccoloVKBuffer>(device, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, sizeof(FrameUBO), nullptr, true);

        // 材质的DescriptorSet的数量需要动态扩容创建
        RecreateMaterialDescriptorPool(NUM_MATERIAL_BATCH);
    }

    void PUnlitMaterialSystem::OnRender(VkCommandBuffer cmdBuffer, PiccoloRenderTarget *renderTarget)
    {
        PScene* scene = GetScene();
        if (!scene) return;

        entt::registry& registry = scene->GetEcsRegistry();
        auto view = registry.view<PTransformComponent, PUnlitMaterialComponent>();
        if (view.begin() == view.end()) return;

        mPipeline->Bind(cmdBuffer);
        PiccoloVKFrameBuffer* frameBuffer = renderTarget->GetCurrentFrameBuffer();
        VkViewport viewport =
        {
            .x = 0,
            .y = 0,
            .width = static_cast<float>(frameBuffer->GetWidth()),
            .height = static_cast<float>(frameBuffer->GetHeight()),
            .minDepth = 0.f,
            .maxDepth = 1.f
        };
        vkCmdSetViewport(cmdBuffer, 0, 1, &viewport);
        VkRect2D scissor =
        {
            .offset = { 0, 0 },
            .extent = { frameBuffer->GetWidth(), frameBuffer->GetHeight() }
        };
        vkCmdSetScissor(cmdBuffer, 0, 1, &scissor);

        // 更新UBO的参数
        UpdateFrameUBODescriptorSet(renderTarget);

        // 检查动态扩容
        bool bShouldForceUpdateMaterial = false;
        uint32_t materialCount = PMaterialFactory::GetInstance()->GetMaterialSize<PUnlitMaterial>();
        if (materialCount > mLastDescriptorSetCount)
        {
            RecreateMaterialDescriptorPool(materialCount);
            bShouldForceUpdateMaterial = true;
        }

        std::vector<bool> updateFlags(materialCount);
        view.each([this, &updateFlags, &bShouldForceUpdateMaterial, &cmdBuffer](const PTransformComponent& transformComponent, const PUnlitMaterialComponent& materialComponent)
        {
            for (const auto & entry : materialComponent.GetMeshMaterials())
            {
                PUnlitMaterial* material = entry.first;
                if (!material || material->GetIndex() < 0)
                {
                    LOG_W("TODO: default material or error material ?");
                    continue;
                }

                uint32_t materialIndex = material->GetIndex();
                VkDescriptorSet paramDescriptorSet = mMaterialDescriptorSetList[materialIndex];
                VkDescriptorSet resourceDescriptorSet = mMaterialResourceDescriptorSetList[materialIndex];

                // 材质只有变动后才更新
                if (!updateFlags[materialIndex])
                {
                    if (material->ShouldFlushParams() || bShouldForceUpdateMaterial)
                    {
                        LOG_T("Update material params : {0}", materialIndex);
                        UpdateMaterialParamDescriptorSet(paramDescriptorSet, material);
                        material->FinishFlushParams();
                    }
                    if (material->ShouldFlushResource() || bShouldForceUpdateMaterial)
                    {
                        LOG_T("Update material resource : {0}", materialIndex);
                        UpdateMaterialResourceDescriptorSet(resourceDescriptorSet, material);
                        material->FinishFlushResource();
                    }
                    updateFlags[materialIndex] = true;
                }

                // 绑定DescriptorSet
                VkDescriptorSet descriptorSets[] = { mFrameUBODescriptorSet, paramDescriptorSet, resourceDescriptorSet };
                vkCmdBindDescriptorSets(cmdBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipelineLayout->GetPipelineLayout(), 0, ARRAY_SIZE(descriptorSets), descriptorSets, 0, nullptr);
                // 推送常量
                ModelPushConstant pushConstant = { transformComponent.GetTransformMatrix() };
                vkCmdPushConstants(cmdBuffer, mPipelineLayout->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(pushConstant), &pushConstant);

                // 绘制使用到这个材质的所有模型
                for (const auto& meshIndex : entry.second)
                {
                    materialComponent.GetMesh(meshIndex)->Draw(cmdBuffer);
                }
            }
        });
    }

    void PUnlitMaterialSystem::OnDestroy()
    {
    }

    void PUnlitMaterialSystem::RecreateMaterialDescriptorPool(uint32_t materialCount)
    {
        PiccoloVKDevice* device = GetDevice();

        uint32_t newDescriptorSetCount = mLastDescriptorSetCount;
        if (mLastDescriptorSetCount == 0)
        {
            newDescriptorSetCount = NUM_MATERIAL_BATCH;
        }

        // 2倍扩容
        while (newDescriptorSetCount < materialCount)
        {
            newDescriptorSetCount <<= 1;
        }
        if (newDescriptorSetCount > NUM_MATERIAL_BATCH_MAX)
        {
            LOG_E("Descriptor Set max count is: {0}, but request: {1}", NUM_MATERIAL_BATCH_MAX, newDescriptorSetCount);
            return;
        }

        LOG_W("{0}: {1} -> {2} S.", __FUNCTION__, mLastDescriptorSetCount, newDescriptorSetCount);

        // 销毁旧的
        mMaterialDescriptorSetList.clear();
        mMaterialResourceDescriptorSetList.clear();
        if (mMaterialDescriptorPool)
        {
            mMaterialDescriptorPool.reset();
        }

        // 重新创建新的pool
        std::vector<VkDescriptorPoolSize> poolSizes =
        {
            {
                .type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                .descriptorCount = newDescriptorSetCount
            },
            {
                .type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
                .descriptorCount = newDescriptorSetCount << 1               // color_tex0、color_tex1
            }
        };
        mMaterialDescriptorPool = std::make_shared<PiccoloVKDescriptorPool>(device, newDescriptorSetCount << 1, poolSizes);         // params、resource

        // 申请分配指定DescriptorSetLayout，指定数量下的DescriptorSet
        mMaterialDescriptorSetList = mMaterialDescriptorPool->AllocateDescriptorSet(mMaterialParamDescriptorSetLayout.get(), newDescriptorSetCount);
        mMaterialResourceDescriptorSetList = mMaterialDescriptorPool->AllocateDescriptorSet(mMaterialResourceDescriptorSetLayout.get(), newDescriptorSetCount);

        assert(mMaterialDescriptorSetList.size() == newDescriptorSetCount && "Failed to AllocateDescriptorSet");
        assert(mMaterialResourceDescriptorSetList.size() == newDescriptorSetCount && "Failed to AllocateDescriptorSet");

        // 创建新的Buffer
        uint32_t diffCount = newDescriptorSetCount - mLastDescriptorSetCount;

        for (int i = 0; i < diffCount; ++i)
        {
            mMaterialBuffers.push_back(std::make_shared<PiccoloVKBuffer>(device, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, sizeof(UnlitMaterialUBO), nullptr, true));
        }

        mLastDescriptorSetCount = newDescriptorSetCount;
    }

    void PUnlitMaterialSystem::UpdateFrameUBODescriptorSet(PiccoloRenderTarget *renderTarget)
    {
        PiccoloApplication* app = GetApplication();
        PiccoloVKDevice* device = GetDevice();

        PiccoloVKFrameBuffer* frameBuffer = renderTarget->GetCurrentFrameBuffer();
        glm::ivec2 resolution = { frameBuffer->GetWidth(), frameBuffer->GetHeight() };
        FrameUBO frameUBO =
        {
            .projMatrix = GetProjMatrix(renderTarget),
            .viewMatrix = GetViewMatrix(renderTarget),
            .resolution = resolution,
            .frameId = static_cast<uint32_t>(app->GetFrameIndex()),
            .time = app->GetStartTimeSecond()
        };
        // 把数据写到Buffer中
        mFrameUBOBuffer->WriteData(&frameUBO);
        // 更新DescriptorSet
        VkDescriptorBufferInfo bufferInfo = DescriptorSetWriter::BuildBufferInfo(mFrameUBOBuffer->GetBuffer(), 0, sizeof(frameUBO));
        VkWriteDescriptorSet bufferWrite = DescriptorSetWriter::WriteBuffer(mFrameUBODescriptorSet, 0, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, &bufferInfo);
        DescriptorSetWriter::UpdateDescriptorSets(device->GetLogicDevice(), { bufferWrite });
    }

    void PUnlitMaterialSystem::UpdateMaterialParamDescriptorSet(VkDescriptorSet descriptorSet, PUnlitMaterial *unlitMaterial)
    {
        PiccoloVKDevice* device = GetDevice();

        PiccoloVKBuffer* materialBuffer = mMaterialBuffers[unlitMaterial->GetIndex()].get();
        UnlitMaterialUBO materialUBO = unlitMaterial->GetMaterialUBO();

        if (const TextureView* texture0 = unlitMaterial->GetTextureView(UNLIT_MAT_BASE_COLOR_0))
        {
            PMaterial::UpdateTextureParams(texture0, &materialUBO.textureParam0);
        }
        if (const TextureView* texture1 = unlitMaterial->GetTextureView(UNLIT_MAT_BASE_COLOR_1))
        {
            PMaterial::UpdateTextureParams(texture1, &materialUBO.textureParam1);
        }

        materialBuffer->WriteData(&materialUBO);
        // 更新DescriptorSet
        VkDescriptorBufferInfo bufferInfo = DescriptorSetWriter::BuildBufferInfo(materialBuffer->GetBuffer(), 0, sizeof(materialUBO));
        VkWriteDescriptorSet bufferWrite = DescriptorSetWriter::WriteBuffer(descriptorSet, 0, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, &bufferInfo);
        DescriptorSetWriter::UpdateDescriptorSets(device->GetLogicDevice(), { bufferWrite });
    }

    void PUnlitMaterialSystem::UpdateMaterialResourceDescriptorSet(VkDescriptorSet descriptorSet, PUnlitMaterial *unlitMaterial)
    {
        PiccoloVKDevice* device = GetDevice();

        const TextureView* textureView0 = unlitMaterial->GetTextureView(UNLIT_MAT_BASE_COLOR_0);
        const TextureView* textureView1 = unlitMaterial->GetTextureView(UNLIT_MAT_BASE_COLOR_1);

        VkDescriptorImageInfo textureInfo0 = DescriptorSetWriter::BuildImageInfo(textureView0->sampler->GetSampler(), textureView0->texture->GetImageView()->GetImageView());
        VkDescriptorImageInfo textureInfo1 = DescriptorSetWriter::BuildImageInfo(textureView1->sampler->GetSampler(), textureView1->texture->GetImageView()->GetImageView());

        VkWriteDescriptorSet textureWrite0 = DescriptorSetWriter::WriteImage(descriptorSet, 0, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, &textureInfo0);
        VkWriteDescriptorSet textureWrite1 = DescriptorSetWriter::WriteImage(descriptorSet, 1, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, &textureInfo1);

        DescriptorSetWriter::UpdateDescriptorSets(device->GetLogicDevice(), { textureWrite0, textureWrite1 });
    }
}
